"""
nums长度小于k，直接返回nums

使用一个heap来保存目前最大的k个元素。当heap中元素个数小于k时，不断向其中加入元素。
"""

import heapq


def small_k_number(nums, k):
    if len(nums) < k:
        return nums
    heap = []
    for i in range(len(nums)):
        if i < k:
            heap.append(nums[i])
        else:
            heapq.heapify(heap)
            if heap[0] < nums[i]:
                heapq.heapreplace(heap, nums[i])
                # Pop and return the current smallest value, and add the new item.
    return heap


print(small_k_number([10, 2, 12, 3, 5, 7], 4))
